package org.wenshu.ai.modular.chat.provider.core;

import org.wenshu.ai.modular.chat.provider.core.chroma.ChromaEmbeddingStoreFactory;
import org.wenshu.ai.modular.chat.provider.core.milvus.MilvusEmbeddingStoreFactory;
import org.wenshu.ai.modular.chat.provider.core.opensearch.OpenSearchEmbeddingStoreFactory;
import org.wenshu.ai.modular.chat.provider.core.pgvector.PgvectorEmbeddingStoreFactory;
import org.wenshu.ai.modular.chat.provider.pojo.EmbeddingStoreConfig;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class EmbeddingStoreFactoryProvider {
    protected static final Map<EmbeddingStoreConfig, EmbeddingStoreFactory> factoryMap =
            new ConcurrentHashMap<>();



    public static EmbeddingStoreFactory getFactory(EmbeddingStoreConfig embeddingStoreConfig) {

        if (EmbeddingStoreType.CHROMA.name().equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
            return factoryMap.computeIfAbsent(embeddingStoreConfig,
                    storeConfig -> new ChromaEmbeddingStoreFactory(storeConfig));
        }
        if (EmbeddingStoreType.MILVUS.name().equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
            return factoryMap.computeIfAbsent(embeddingStoreConfig,
                    storeConfig -> new MilvusEmbeddingStoreFactory(storeConfig));
        }
        if (EmbeddingStoreType.PGVECTOR.name()
                .equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
            return factoryMap.computeIfAbsent(embeddingStoreConfig,
                    storeConfig -> new PgvectorEmbeddingStoreFactory(storeConfig));
        }

        if (EmbeddingStoreType.OPENSEARCH.name()
                .equalsIgnoreCase(embeddingStoreConfig.getProvider())) {
            return factoryMap.computeIfAbsent(embeddingStoreConfig,
                    storeConfig -> new OpenSearchEmbeddingStoreFactory(storeConfig));
        }
        throw new RuntimeException("Unsupported EmbeddingStoreFactory provider: "
                + embeddingStoreConfig.getProvider());
    }
}
